import torch
from torch import nn
from src.model.time_model.gru import Gru
from src.model.time_model.tcn import Tcn
from mamba_ssm import Mamba


class TimeBlock(nn.Module):
    def __init__(self, dims, seq, pred_len, num_nodes, num_layer, single_point=True):
        super(TimeBlock, self).__init__()
        self.num_nodes = num_nodes
        self.num_layer = num_layer
        self.input_size, self.hidden_size, self.output_size = dims
        self.seq = seq
        self.pred_len = pred_len
        self.single_point = single_point
        
        self.fc1 = nn.Linear(self.input_size, self.hidden_size)
        self.mamba = Mamba(d_model=self.hidden_size, d_state=64, d_conv=8, expand=2)
        self.fc2 = nn.Linear(self.hidden_size, self.output_size)

    def forward(self, x):
        # x是(B, T, N, C)
        x = self.fc1(x)
        B, T, N, H = x.shape
        x = x.permute(0, 2, 1, 3).contiguous()
        x = x.reshape(-1, T, H)
        y = self.mamba(x)
        # yi会变成(B * N, T, C)
        y = y.reshape(B, T, N, H)
        y = y[:, -1, :, :]
        return self.fc2(y)